from config import PathVar, parametersNvariables
from model_2S3EIAR import Learner,loss
import matplotlib.pyplot as plt

resultTwo = True
# resultTwo = False


if resultTwo:
    print('second stage')
    ## 拟合
    pathvar = PathVar()
    pv = parametersNvariables()
    city = '绍兴'
    path = pathvar.covid_csv
    print('选择城市:', city)
    print('文件地址:', path)
    learner = Learner(loss, city=city, path=path, parameter_output = pathvar.resultTwoParameters)
    range_parameters = pv.init_range_parameters()
    print('range parameters:', range_parameters)
    start_variables = pv.init_start_variables()
    print('start variables:', start_variables)
    start_parameters = pv.init_start_parameters()
    print('start parameters:', start_parameters)

    learner.input_range_parameters(range_parameters.copy())
    print('input range parameters, done!')
    learner.input_start_variables(start_variables.copy())
    print('input start variables, done!')
    learner.input_start_parameters(start_parameters.copy())
    print('input start parameters, done!')
    df, result = learner.train()
    fig, ax = plt.subplots(figsize=(15, 10))
    df.plot(ax=ax)
    fig.show()
    # fig, ax = plt.subplots(figsize=(15, 10))
    # df.plot(ax=ax)
    # fig.show()
    df.to_excel(pathvar.resultTwoExcel)